#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
图片处理性能测试脚本
测试不同大小图片的处理时间，用于优化自动扩缩容策略
"""

import os
import time
import requests
import json
import base64
from datetime import datetime
from typing import Dict, List, Any
import statistics

class ImageProcessingPerformanceTest:
    """图片处理性能测试器"""
    
    def __init__(self, api_base_url: str = "http://localhost:8080"):
        self.api_base_url = api_base_url
        self.test_results: List[Dict[str, Any]] = []
        
    def encode_image_to_base64(self, image_path: str) -> str:
        """将图片编码为base64"""
        with open(image_path, 'rb') as image_file:
            encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
            return f"data:image/jpeg;base64,{encoded_string}"
    
    def get_file_size_mb(self, file_path: str) -> float:
        """获取文件大小(MB)"""
        size_bytes = os.path.getsize(file_path)
        return size_bytes / (1024 * 1024)
    
    def test_single_image(self, image_path: str, test_name: str) -> Dict[str, Any]:
        """测试单张图片处理性能"""
        print(f"🧪 测试图片: {test_name}")
        print(f"📁 文件路径: {image_path}")
        
        # 获取文件信息
        file_size_mb = self.get_file_size_mb(image_path)
        print(f"📊 文件大小: {file_size_mb:.2f} MB")
        
        # 编码图片
        print("🔄 编码图片为base64...")
        encode_start = time.time()
        base64_image = self.encode_image_to_base64(image_path)
        encode_time = time.time() - encode_start
        print(f"✅ 编码完成: {encode_time:.2f}秒")
        
        # 准备请求数据
        image_data = base64_image.split(',')[1]  # 移除data:image/jpeg;base64,前缀
        image_bytes = base64.b64decode(image_data)
        
        # 获取可用的真实应用
        print("🔍 获取可用的真实应用...")
        app_response = requests.get(f"{self.api_base_url}/api/v1/real-apps/next")
        if app_response.status_code != 200:
            print(f"❌ 获取应用失败: {app_response.status_code}")
            return None
        
        app_data = app_response.json()
        if not app_data.get('success'):
            print(f"❌ 获取应用失败: {app_data.get('error')}")
            return None
        
        app_id = app_data['data']['app_id']
        print(f"✅ 获取到应用: {app_id}")
        
        # 发送处理请求到真实应用
        print("🚀 发送处理请求到真实应用...")
        process_start = time.time()
        
        try:
            # 准备任务数据
            task_data = {
                'task_id': f"test_{int(time.time())}",
                'image_data': base64_image,
                'tile_size': '400',
                'quality_level': 'high'
            }
            
            response = requests.post(
                f"{self.api_base_url}/api/v1/real-apps/{app_id}/task",
                json=task_data,
                timeout=300  # 5分钟超时
            )
            
            process_end = time.time()
            processing_time = process_end - process_start
            
            if response.status_code == 200:
                result = response.json()
                print(f"✅ 处理成功: {processing_time:.2f}秒")
                
                # 检查是否是异步任务
                if result.get('task_id'):
                    print("🔄 检测到异步任务，开始轮询结果...")
                    final_result = self.poll_task_result(result['task_id'])
                    if final_result:
                        processing_time = final_result['total_time']
                        print(f"✅ 异步处理完成: {processing_time:.2f}秒")
                    else:
                        print("❌ 异步任务轮询失败")
                        return None
                
                return {
                    'test_name': test_name,
                    'file_path': image_path,
                    'file_size_mb': file_size_mb,
                    'encode_time': encode_time,
                    'processing_time': processing_time,
                    'total_time': encode_time + processing_time,
                    'success': True,
                    'timestamp': datetime.now().isoformat()
                }
            else:
                print(f"❌ 处理失败: {response.status_code}")
                return {
                    'test_name': test_name,
                    'file_path': image_path,
                    'file_size_mb': file_size_mb,
                    'encode_time': encode_time,
                    'processing_time': 0,
                    'total_time': encode_time,
                    'success': False,
                    'error': f"HTTP {response.status_code}",
                    'timestamp': datetime.now().isoformat()
                }
                
        except requests.exceptions.Timeout:
            print("⏰ 请求超时")
            return {
                'test_name': test_name,
                'file_path': image_path,
                'file_size_mb': file_size_mb,
                'encode_time': encode_time,
                'processing_time': 0,
                'total_time': encode_time,
                'success': False,
                'error': 'Request timeout',
                'timestamp': datetime.now().isoformat()
            }
        except Exception as e:
            print(f"❌ 请求异常: {e}")
            return {
                'test_name': test_name,
                'file_path': image_path,
                'file_size_mb': file_size_mb,
                'encode_time': encode_time,
                'processing_time': 0,
                'total_time': encode_time,
                'success': False,
                'error': str(e),
                'timestamp': datetime.now().isoformat()
            }
    
    def poll_task_result(self, task_id: str, max_attempts: int = 60, interval: int = 5) -> Dict[str, Any]:
        """轮询异步任务结果"""
        print(f"🔍 轮询任务: {task_id}")
        
        start_time = time.time()
        
        for attempt in range(1, max_attempts + 1):
            try:
                response = requests.get(
                    f"{self.api_base_url}/api/v1/tasks/{task_id}/status",
                    timeout=10
                )
                
                if response.status_code == 200:
                    result = response.json()
                    status = result.get('status', 'unknown')
                    
                    print(f"📊 第{attempt}次查询: {status}")
                    
                    if status == 'completed':
                        total_time = time.time() - start_time
                        return {
                            'total_time': total_time,
                            'result': result
                        }
                    elif status == 'failed':
                        print(f"❌ 任务失败: {result.get('error', 'Unknown error')}")
                        return None
                    elif status in ['processing', 'queued']:
                        time.sleep(interval)
                        continue
                    else:
                        print(f"⚠️ 未知状态: {status}")
                        time.sleep(interval)
                        continue
                else:
                    print(f"❌ 状态查询失败: {response.status_code}")
                    time.sleep(interval)
                    continue
                    
            except Exception as e:
                print(f"❌ 轮询异常: {e}")
                time.sleep(interval)
                continue
        
        print("⏰ 轮询超时")
        return None
    
    def run_batch_test(self, test_images_dir: str) -> List[Dict[str, Any]]:
        """运行批量测试"""
        print("🚀 开始图片处理性能测试")
        print("=" * 60)
        
        if not os.path.exists(test_images_dir):
            print(f"❌ 测试目录不存在: {test_images_dir}")
            return []
        
        # 获取所有测试图片
        test_images = []
        for filename in os.listdir(test_images_dir):
            if filename.lower().endswith(('.jpg', '.jpeg')):
                file_path = os.path.join(test_images_dir, filename)
                test_images.append((file_path, filename))
        
        if not test_images:
            print("❌ 没有找到测试图片")
            return []
        
        # 按文件大小排序
        test_images.sort(key=lambda x: os.path.getsize(x[0]))
        
        print(f"📁 找到 {len(test_images)} 张测试图片")
        print()
        
        # 逐个测试
        for i, (file_path, filename) in enumerate(test_images, 1):
            print(f"📸 测试 {i}/{len(test_images)}: {filename}")
            result = self.test_single_image(file_path, filename)
            if result:
                self.test_results.append(result)
            print("-" * 40)
            print()
        
        return self.test_results
    
    def analyze_results(self) -> Dict[str, Any]:
        """分析测试结果"""
        if not self.test_results:
            return {}
        
        successful_results = [r for r in self.test_results if r['success']]
        failed_results = [r for r in self.test_results if not r['success']]
        
        if not successful_results:
            return {
                'total_tests': len(self.test_results),
                'successful_tests': 0,
                'failed_tests': len(failed_results),
                'success_rate': 0,
                'error': 'No successful tests'
            }
        
        # 计算统计信息
        file_sizes = [r['file_size_mb'] for r in successful_results]
        processing_times = [r['processing_time'] for r in successful_results]
        total_times = [r['total_time'] for r in successful_results]
        
        analysis = {
            'total_tests': len(self.test_results),
            'successful_tests': len(successful_results),
            'failed_tests': len(failed_results),
            'success_rate': len(successful_results) / len(self.test_results) * 100,
            
            'file_size_stats': {
                'min_mb': min(file_sizes),
                'max_mb': max(file_sizes),
                'avg_mb': statistics.mean(file_sizes),
                'median_mb': statistics.median(file_sizes)
            },
            
            'processing_time_stats': {
                'min_seconds': min(processing_times),
                'max_seconds': max(processing_times),
                'avg_seconds': statistics.mean(processing_times),
                'median_seconds': statistics.median(processing_times)
            },
            
            'total_time_stats': {
                'min_seconds': min(total_times),
                'max_seconds': max(total_times),
                'avg_seconds': statistics.mean(total_times),
                'median_seconds': statistics.median(total_times)
            },
            
            'performance_by_size': self._analyze_performance_by_size(successful_results),
            'failed_tests': failed_results
        }
        
        return analysis
    
    def _analyze_performance_by_size(self, results: List[Dict[str, Any]]) -> Dict[str, Any]:
        """按文件大小分析性能"""
        size_categories = {
            'small': [],    # < 200KB
            'medium': [],   # 200KB - 500KB
            'large': []     # > 500KB
        }
        
        for result in results:
            size_mb = result['file_size_mb']
            if size_mb < 0.2:
                size_categories['small'].append(result)
            elif size_mb < 0.5:
                size_categories['medium'].append(result)
            else:
                size_categories['large'].append(result)
        
        analysis = {}
        for category, category_results in size_categories.items():
            if category_results:
                processing_times = [r['processing_time'] for r in category_results]
                analysis[category] = {
                    'count': len(category_results),
                    'avg_processing_time': statistics.mean(processing_times),
                    'min_processing_time': min(processing_times),
                    'max_processing_time': max(processing_times)
                }
        
        return analysis
    
    def save_results(self, output_file: str = None) -> str:
        """保存测试结果"""
        if output_file is None:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            output_file = f"image_processing_test_results_{timestamp}.json"
        
        results = {
            'test_info': {
                'timestamp': datetime.now().isoformat(),
                'api_base_url': self.api_base_url,
                'total_tests': len(self.test_results)
            },
            'raw_results': self.test_results,
            'analysis': self.analyze_results()
        }
        
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(results, f, indent=2, ensure_ascii=False)
        
        print(f"💾 测试结果已保存到: {output_file}")
        return output_file
    
    def print_summary(self):
        """打印测试摘要"""
        analysis = self.analyze_results()
        
        print("📊 测试结果摘要")
        print("=" * 60)
        print(f"总测试数: {analysis.get('total_tests', 0)}")
        print(f"成功测试: {analysis.get('successful_tests', 0)}")
        print(f"失败测试: {analysis.get('failed_tests', 0)}")
        print(f"成功率: {analysis.get('success_rate', 0):.1f}%")
        print()
        
        if analysis.get('successful_tests', 0) > 0:
            print("📈 性能统计:")
            print(f"  文件大小范围: {analysis['file_size_stats']['min_mb']:.2f} - {analysis['file_size_stats']['max_mb']:.2f} MB")
            print(f"  平均文件大小: {analysis['file_size_stats']['avg_mb']:.2f} MB")
            print()
            print(f"  处理时间范围: {analysis['processing_time_stats']['min_seconds']:.1f} - {analysis['processing_time_stats']['max_seconds']:.1f} 秒")
            print(f"  平均处理时间: {analysis['processing_time_stats']['avg_seconds']:.1f} 秒")
            print(f"  中位数处理时间: {analysis['processing_time_stats']['median_seconds']:.1f} 秒")
            print()
            print(f"  总时间范围: {analysis['total_time_stats']['min_seconds']:.1f} - {analysis['total_time_stats']['max_seconds']:.1f} 秒")
            print(f"  平均总时间: {analysis['total_time_stats']['avg_seconds']:.1f} 秒")
            print()
            
            # 按大小分类的性能
            size_perf = analysis.get('performance_by_size', {})
            if size_perf:
                print("📊 按文件大小分类的性能:")
                for size_category, perf in size_perf.items():
                    print(f"  {size_category}: {perf['count']}张图片, 平均{perf['avg_processing_time']:.1f}秒")


def main():
    """主函数"""
    # 测试图片目录
    test_images_dir = "/home/ubuntu/PhotoEnhanceAI-web/test-jpg"
    
    # API基础URL
    api_base_url = "http://localhost:8080"
    
    # 创建测试器
    tester = ImageProcessingPerformanceTest(api_base_url)
    
    # 运行测试
    results = tester.run_batch_test(test_images_dir)
    
    # 打印摘要
    tester.print_summary()
    
    # 保存结果
    output_file = tester.save_results()
    
    print(f"\n🎯 测试完成！结果已保存到: {output_file}")


if __name__ == "__main__":
    main()
